import torch
from torch import nn
from layers.dummy_layers import DummyLayer

class DummyNet(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1):
        super(DummyNet, self).__init__()
        self.dummy_layer = DummyLayer(in_channels,out_channels,kernel_size,stride)


    def forward(self,x):
        return self.dummy_layer(x)


def get_dummy_net(input_shape):
    return DummyNet(
        input_shape[1],
        input_shape[1],
    )


if __name__ == '__main__':
    batch_shape = (32, 3, 224, 224)
    fake_batch = torch.randn(batch_shape)
    nets = [
        get_dummy_net(batch_shape)
    ]
    for Net in nets:
        output = Net(fake_batch)
        print(type(Net), 'output shape:', output.shape)
